import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import re
import matplotlib.colors as mcolors
# 设置英文和数字的字体为 Times New Roman
plt.rcParams['font.family'] = 'Times New Roman'
file_path = 'sum_time.xlsx'
xls = pd.ExcelFile(file_path)
# 获取所有子表的名称
sheet_names = xls.sheet_names
# 定义要处理的分布模式
distribution_modes = ['normal_crown', 'skew_normal_crown', 'uniform_crown']
# 定义一个函数来清理 clade 列中的中文字符和括号
def clean_clade_name(clade_name):
    # 使用正则表达式移除中文字符和括号
    cleaned_name = re.sub(r'[\u4e00-\u9fff()（）]', '', clade_name)
    return cleaned_name
# 为每种分布模式创建一个图表
for mode in distribution_modes:
    # 合并所有子表的数据
    combined_data = []
    for sheet_name in sheet_names:
        df = pd.read_excel(xls, sheet_name=sheet_name)
        if mode in df.columns:
            # 提取物种名称和当前分布模式的数据并清理 clade 列中非必要字符
            mode_data = df[['clade', mode]].copy()
            mode_data['clade'] = mode_data['clade'].apply(clean_clade_name)
            mode_data['source'] = sheet_name  # 添加来源列
            combined_data.append(mode_data)
    if combined_data:
        combined_df = pd.concat(combined_data, ignore_index=True)
        # 数据清理：移除包含非数值的行
        combined_df = combined_df[combined_df[mode].apply(lambda x: isinstance(x, str) and '-' in x)]
        all_sources = combined_df['source'].unique()
        # 为每个来源分配一个唯一的颜色
        source_colors = {}
        color_map = plt.cm.tab20  # 使用 tab20 色彩映射，支持更多颜色
        for i, source in enumerate(all_sources):
            source_colors[source] = color_map(i % 20)
        # 提取物种名称列表并为每种分布模式创建一个图表
        species_list = combined_df['clade'].unique()
        plt.figure(figsize=(15, 10))
        # 计算所有中值以确定 y 轴范围
        all_median_times = []
        for species in species_list:
            species_data = combined_df[combined_df['clade'] == species]
            if not species_data.empty:
                time_ranges = species_data[mode].str.split('-', expand=True).astype(float)
                median_time = time_ranges.apply(lambda x: (x[0] + x[1]) / 2, axis=1).values
                all_median_times.extend(median_time)
        # 设置 y 轴范围并为每个物种绘制柱状图
        min_time = np.floor(min(all_median_times))
        max_time = np.ceil(max(all_median_times))
        plt.ylim(min_time, max_time)
        plt.yticks(np.arange(min_time, max_time + 1, 1))
        bar_width = 0.2
        index = np.arange(len(species_list))
        handles = []
        labels = []
        for i, species in enumerate(species_list):
            species_data = combined_df[combined_df['clade'] == species]
            if not species_data.empty:
                # 计算每个区间的中值
                time_ranges = species_data[mode].str.split('-', expand=True).astype(float)
                median_time = time_ranges.apply(lambda x: (x[0] + x[1]) / 2, axis=1).values
                # 获取来源列，绘制柱状图
                sources = species_data['source'].values
                positions = index[i] + np.arange(len(median_time)) * bar_width
                for j, (time, source) in enumerate(zip(median_time, sources)):
                    color = source_colors.get(source, 'gray')
                    bar = plt.bar(positions[j], time, width=bar_width, color=color, label=source)
                    if source not in labels:
                        handles.append(bar)
                        labels.append(source)
        plt.title(f'Time of Differentiation of Species Crown Groups ({mode})')
        plt.xlabel('Species')
        plt.ylabel('Differentiation Time (billions of years)')
        plt.xticks(index + bar_width * 1.5, species_list, rotation=45, ha='right')
        plt.legend(handles=handles, labels=labels, title='Source')
        plt.tight_layout()
        plt.savefig(f'{mode}_differentiation_time.png', dpi=300)
        plt.close()
    else:
        print(f"No data to plot for {mode}.")